Skip to content

Reindex via indexer#1032

Open
sharkinsspatial wants to merge 7 commits into
mainfrom
reindex-via-indexer
Open

Reindex via indexer#1032
sharkinsspatial wants to merge 7 commits into
mainfrom
reindex-via-indexer

Conversation

@sharkinsspatial

@sharkinsspatial sharkinsspatial commented Jun 30, 2026

Copy link
Copy Markdown
Collaborator

What I did

Thanks to @dcherian for recommending this approach 🙇

Added support for ManifestArray to handle xarray's reindex/alignment indexer at the array layer, so vanilla xr.reindex, xr.align(join="outer") that would introduce nans work with virtual arrays and stay lazy. Missing positions become null-path manifest entries that Zarr returns as the array's fill_value — no materialization.

  • indexing.py — route a "fancy" integer-array indexer to a chunk-grid remap instead of rejecting it, _collapse_outer_indexer reshapes broadcast (vectorized) indexers back to their 1D per-axis form.
  • reindex.py (new) — chunk_map_from_indexer translates xarray's per-element reindex indexer (e.g. [0,1,2,-1,-1], where -1 marks a missing label) into a per-chunk plan. It slices the indexer into chunk-sized blocks and requires each block to be either all -1 → emit an empty null-path chunk, or a contiguous run that exactly fills one source chunk → reuse that chunk's reference. Anything else (a block mixing real and missing positions, or a partial/reordered chunk) would mean materializing, so it raises. The result is one source-chunk index or null chunk per target chunk slot.
  • array_api.py — np.where returns the gathered array unchanged when the fill mask matches its null chunks.
  • array.py — adds ManifestArray._reindex_axis and admits ndarray in __array_function__.

This replaces the bespoke manifest surgery people need today to align virtual datasets onto a common grid when the alignment would result in nans.

By design there's no chunk splitting / within-chunk reordering — target labels must land on chunk boundaries. Because align sorts the union ascending, a descending coordinate (e.g. ITS_LIVE y coord) requires explicit reindex with a descending target.

Teach ManifestArray to handle xarray's reindex/alignment indexer at the
array-protocol layer, so plain xr.reindex, xr.align, and xr.concat work over
virtual arrays and stay lazy — missing positions become null-path manifest
entries that read back as the array's fill_value, with no materialization and
no dtype promotion.

xarray's Variable._getitem_with_mask gathers with an integer indexer (-1 for
missing) then masks with where(). The hooks that make this work lazily:

- manifests/indexing.py: an integer-array indexer is routed to a chunk-grid
  remap (chunk_map_from_indexer -> ManifestArray._reindex_axis) instead of being
  rejected as fancy indexing. Each chunk-block must be all-missing (-> null
  chunk) or a chunk-aligned run (-> real chunk), else NotImplementedError; no
  chunk splitting.
- manifests/reindex.py (new): chunk_map_from_indexer partitions the indexer by
  chunk and asserts the all-or-nothing condition.
- array_api.py: np.where returns the gathered array unchanged when the requested
  fill mask matches its null chunks (the manifest's own fill_value governs);
  general masking still raises. np.result_type treats a lone ManifestArray's
  dtype as authoritative, so int arrays keep their dtype + declared sentinel fill
  rather than being promoted to float+NaN (also fixes a latent
  generator-consumption bug in the previous implementation).
- array.py: __array_function__ now admits ndarray among the argument types
  (np.where's boolean condition arrives as one); adds ManifestArray._reindex_axis.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@codecov

codecov Bot commented Jun 30, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.18182% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.99%. Comparing base (e1ebb4b) to head (c82094f).

Files with missing lines Patch % Lines
virtualizarr/manifests/array.py 87.87% 4 Missing ⚠️
virtualizarr/manifests/indexing.py 84.61% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1032      +/-   ##
==========================================
+ Coverage   89.87%   89.99%   +0.11%     
==========================================
  Files          36       37       +1     
  Lines        2202     2288      +86     
==========================================
+ Hits         1979     2059      +80     
- Misses        223      229       +6     
Files with missing lines Coverage Δ
virtualizarr/manifests/array_api.py 98.61% <100.00%> (+0.42%) ⬆️
virtualizarr/manifests/reindex.py 100.00% <100.00%> (ø)
virtualizarr/manifests/indexing.py 92.42% <84.61%> (-1.92%) ⬇️
virtualizarr/manifests/array.py 85.29% <87.87%> (+0.67%) ⬆️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

A 2D+ outer-join reindexes more than one dimension in a single call, and
xarray expresses that as a broadcast (vectorized) indexer: one integer array
per axis, shaped to the array's rank with the real labels on its own axis and
length 1 everywhere else (e.g. (N,1) and (1,M)). _apply_reindex assumed a 1D
indexer per axis, so it misread the array's rank — raising "axis 0 (chunk
size 1)" or, on newer numpy, TypeError from int() on a non-scalar block.

Add _collapse_outer_indexer: since reindex indexing is orthogonal, each
broadcast component reshapes losslessly to its 1D per-axis indexer. A genuine
pointwise/vectorized indexer (non-axis dim != 1) is rejected with a clear
error. This makes the real ITS_LIVE x+y mosaic work via native reindex/concat.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

@TomNicholas TomNicholas left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Excited to get this to work, but I think we'll have to go back and forth re-organizing and polishing before I feel confident about this.

Comment thread virtualizarr/tests/test_manifests/test_reindex.py
Comment thread virtualizarr/tests/test_native_align.py Outdated
Comment thread virtualizarr/tests/test_native_align.py Outdated
Comment thread virtualizarr/tests/test_native_align.py Outdated
@TomNicholas

Copy link
Copy Markdown
Member

Also it would be quite useful to have an issue/docs example / both showing how this actually solves user's NaN-padding and alignment problems.

@sharkinsspatial

Copy link
Copy Markdown
Collaborator Author

Tests are failing due to an upstream GRIB parser issue.

@sharkinsspatial sharkinsspatial marked this pull request as ready for review July 9, 2026 17:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support for sparsely tiled spatial grids?

2 participants